/*
 *actually 3d-partitioning in this program 
 */


#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <assert.h>
#include <sys/time.h>
#include "3d7p_aceMesh.h"
#include "aceMesh_runtime_c.h"

double test[2][SIZEX][SIZEY][SIZEZ];
double A[2][SIZEX][SIZEY][SIZEZ];

extern void  slave_stencil_core(targs* args);

#define alpha_d 0.0876
#define beta_d  0.0765

#define TestDo3d7pkernel(itr,i,j,k) {\
   test[itr%2][i][j][k] = alpha_d * test[(itr+1)%2][i][j][k] + \
                        beta_d * (test[(itr+1)%2][i-1][j][k] + test[(itr+1)%2][i+1][j][k] +\
                        test[(itr+1)%2][i][j-1][k] + test[(itr+1)%2][i][j+1][k] +\
                        test[(itr+1)%2][i][j][k-1] + test[(itr+1)%2][i][j][k+1] ); } 

void my_neighbors(void** neighbor_addrs,int* num_neigh, targs* args)
{
  int starti = args->starti;
  int startj = args->startj;
  int startk = args->startk;
  void* src_addr=args->src_addr;
  *num_neigh=0;
  if(starti >= BLKX)
     neighbor_addrs[(*num_neigh)++] = (void*)((unsigned long)src_addr - BLKX * SIZEY * SIZEZ * 8);  
  if(starti < SIZEX - BLKX)
     neighbor_addrs[(*num_neigh)++] = (void*)((unsigned long)src_addr + BLKX * SIZEY * SIZEZ * 8);
  if(startj >= BLKY)
     neighbor_addrs[(*num_neigh)++] = ((void*)((unsigned long)src_addr - BLKY * SIZEZ * 8));
  if(startj < SIZEY - BLKY)
     neighbor_addrs[(*num_neigh)++] = ((void*)((unsigned long)src_addr + BLKY * SIZEZ* 8));
  if(startk >= BLKZ)
     neighbor_addrs[(*num_neigh)++] = ((void*)((unsigned long)src_addr - BLKZ * 8));
  if(startk < SIZEZ - BLKZ)
     neighbor_addrs[(*num_neigh)++] = ((void*)((unsigned long)src_addr + BLKZ * 8));
}


int Do3d7p( const int iter, const int thread_num)
{
  targs each;
  int i, j, k, itr;
  
  for( itr = 0; itr < iter; ++itr)
  {
     acemesh_begin_split_task("3d7p iter"); 
     for( i = 1; i < SIZEX -1; i += BLKX)
     {
        int endi = (i + BLKX > SIZEX -1 ) ? SIZEX - 1  : i + BLKX;
        for( j = 1; j < SIZEY -1; j += BLKY)
        {
           int endj = (j + BLKY > SIZEY -1 ) ? SIZEY - 1 : j + BLKY;
           for( k = 1; k < SIZEZ -1; k += BLKZ)
           {
              int endk = (k + BLKZ > SIZEZ -1 ) ? SIZEZ - 1 : k + BLKZ;
              each.starti = i;
              each.endi = endi;
              each.startj = j;
              each.endj = endj;
              each.startk = k;
              each.endk = endk;
              each.iter = itr;
              each.src_addr=&A[(itr+1)%2][i][j][k];
              acemesh_push_wlist(1, &A[itr%2][i][j][k],NORMAL);
              acemesh_push_rlist(1, &A[(itr+1)%2][i][j][k], NORMAL);
              acemesh_task_generator_with_neighbors((TASK_FUNCPTR)slave_stencil_core,
                                                    (void*)(&each), sizeof(targs), 
                                                    NULL, (NEIGHBOR_FUNCPTR)my_neighbors, &each);
              acemesh_task_set_type(STENCIL_TASK);
           }  
        }
     }
     acemesh_end_split_task();
  }
  acemesh_spawn_and_wait(1);
  return 0;
}


void init()
{
    int i,j,k;
    for( i = 0; i < SIZEX; i++)
     for( j = 0; j < SIZEY; j++)
      for( k = 0; k < SIZEZ; k++){
        A[1][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0; 
     }
}
int check(const int iter)
{
   int i,j,k;
   int itr;
   int ii,jj,kk;
   printf("begin check\n");
    for(i = 0; i < SIZEX; i++)
     for( j = 0; j < SIZEY; j++)
      for( k = 0; k < SIZEZ; k++){
         test[1][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0; 
      }
       
    for(itr = 0; itr < iter; ++itr){
     for(i = 1; i < SIZEX - 1; i++)
      for( j = 1; j < SIZEY - 1; j++)
       for( k = 1; k < SIZEZ - 1; k++)
         TestDo3d7pkernel(itr,i,j,k);
      }
      
    for(ii = 1; ii < SIZEX - 1; ii++)
      for(jj = 1; jj < SIZEY - 1; jj++)
       for(kk = 1; kk < SIZEZ - 1; kk++){
          if( fabs(test[(iter+1)%2][ii][jj][kk]  - A[(iter+1)%2][ii][jj][kk]) > 1e-13 ){
             printf("<i: %d   j:  %d k: %d> ",ii,jj,kk);    
             printf("test: %lf",test[(iter+1)%2][ii][jj][kk]);
             printf(", but A: %lf",A[(iter+1)%2][ii][jj][kk]);
             printf("error\n");               
             return -1;
         }
      }
  printf("correct\n");
  return 0;
}



int main(int argc, char** argv)
{
    int ITER,num_threads;
    int i,j,k,p;
    struct timeval start;
    struct timeval end;
    double total_time;       
    
    ITER = 30;
    num_threads = 16;
    if(argc > 1)
       num_threads = atoi(argv[1]);
    if(argc > 2)
       ITER = atoi(argv[2]);
    if(((SIZEX-2)%BLKX) != 0 || ((SIZEY-2)%BLKY) != 0 || ((SIZEZ-2)%BLKZ) != 0)
    {   
      printf("the expressions SIZEX-2)%BLKX ,(SIZEY-2)%BLKY and  (SIZEZ-2)%BLKZ must == 0  \n");
      assert(0);
    }  
    printf("iter=%d",ITER);
    init();
    p=1;
    //while (p);
    acemesh_runtime_init(num_threads);

    gettimeofday(&start,NULL);
    Do3d7p( ITER, num_threads);
    gettimeofday(&end,NULL);   

    aceMesh_runtime_shutdown();  
    total_time = 0.0;
    total_time+=(end.tv_sec-start.tv_sec)+(end.tv_usec-start.tv_usec)/1000000.0;
#ifdef CHECK
    if(!check(ITER))
    {
        printf("total_time: %lf\n",total_time);       
    }
    else 
    {
        printf("total_time: %lf check_error\n", total_time);
        exit(1);
    }
#else
    printf("total_time: %lf\n",total_time);       
#endif
    return 0;
}
